package com.mmg.cosplay.service;

import com.mmg.cosplay.entity.po.RolePreset;
import dev.langchain4j.data.segment.TextSegment;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 增强聊天服务，集成知识库检索
 * 
 * @author fly
 * @since 2025-9-24
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class EnhancedChatService {

    private final ChatService chatService;
    private final IKnowledgeService knowledgeService;
    private final IRolePresetService rolePresetService;

    /**
     * 带知识库检索的聊天
     * 
     * @param memoryId 会话ID
     * @param roleId 角色ID
     * @param userMessage 用户消息
     * @return 流式响应
     */
    public Flux<String> chatWithKnowledge(String memoryId, Long roleId, String userMessage) {
        try {
            // 1. 获取角色信息
            RolePreset rolePreset = rolePresetService.getById(roleId);
            if (rolePreset == null) {
                return Flux.error(new RuntimeException("角色不存在"));
            }

            // 2. 在角色知识库中搜索相关内容
            List<TextSegment> relevantSegments = knowledgeService.searchKnowledge(roleId, userMessage, 3);
            
            // 3. 构建增强的系统提示词
            String enhancedSystemPrompt = buildEnhancedSystemPrompt(rolePreset.getSystemPrompt(), relevantSegments);
            
            // 4. 调用原始聊天服务
            return chatService.fluxChat(memoryId, enhancedSystemPrompt, userMessage);
            
        } catch (Exception e) {
            log.error("知识库检索聊天失败: memoryId={}, roleId={}, userMessage={}", memoryId, roleId, userMessage, e);
            // 如果知识库检索失败，降级到普通聊天
            RolePreset rolePreset = rolePresetService.getById(roleId);
            return chatService.fluxChat(memoryId, rolePreset.getSystemPrompt(), userMessage);
        }
    }

    /**
     * 普通聊天（不使用知识库）
     */
    public Flux<String> chatWithoutKnowledge(String memoryId, Long roleId, String userMessage) {
        RolePreset rolePreset = rolePresetService.getById(roleId);
        return chatService.fluxChat(memoryId, rolePreset.getSystemPrompt(), userMessage);
    }

    /**
     * 构建增强的系统提示词
     */
    private String buildEnhancedSystemPrompt(String originalPrompt, List<TextSegment> relevantSegments) {
        if (relevantSegments.isEmpty()) {
            return originalPrompt;
        }

        StringBuilder enhancedPrompt = new StringBuilder();
        enhancedPrompt.append(originalPrompt);
        enhancedPrompt.append("\n\n");
        enhancedPrompt.append("=== 相关知识库内容 ===\n");
        enhancedPrompt.append("以下是与用户问题相关的知识库内容，请结合这些信息来回答用户的问题：\n\n");

        for (int i = 0; i < relevantSegments.size(); i++) {
            TextSegment segment = relevantSegments.get(i);
            enhancedPrompt.append("知识片段 ").append(i + 1).append(":\n");
            enhancedPrompt.append(segment.text()).append("\n\n");
        }

        enhancedPrompt.append("=== 知识库内容结束 ===\n\n");
        enhancedPrompt.append("请基于上述知识库内容和你的角色设定来回答用户的问题。如果知识库中没有相关信息，请按照你的角色设定正常回答。");

        return enhancedPrompt.toString();
    }
}